import baidu_ai_common
import requests
import os
import json


def main():
    group_id = "facelib_test"
    update_face_lib = int(input('是否更新人脸库？（1：是，0：否）：'))
    if update_face_lib == 1:
        updateFaceLib(group_id)

    print('---------开始识别“人脸照片”文件夹中的照片----------')
    photo_path = "人脸照片/"
    photo_files = os.listdir(photo_path)

    if len(photo_files) == 0:
        print('提示：文件夹人脸照片中没有任何需要识别的照片，请先放置照片到此文件夹中')
    else:
        url = "https://aip.baidubce.com/rest/2.0/face/v3/search?access_token=" + \
            baidu_ai_common.get_access_token()
        headers = {
            'Content-Type': 'application/json'
        }
        i=1
        for photo in photo_files:
            photo_info = photo.split('.')
            photo_ext = photo_info[1]
            if photo_ext.lower() in ['jpg', 'jpeg', 'png', 'bmp']:
                payload = json.dumps({
                    "group_id_list": group_id,
                    "image": baidu_ai_common.get_file_content_as_base64(photo_path+photo, False),
                    "image_type": "BASE64"
                })
                response = requests.request("POST", url, headers=headers, data=payload)
                result=response.json()
                print('图片：',photo,'的识别结果：')
                if result["error_code"]==0:
                    print('识别成功（',i,'）：',result['result']['user_list'][0]['user_info'])
                else:
                    print(photo,'，识别失败（',i,'）')
                print('----------------------------------')
                i+=1
        print('文件夹：“人脸照片”中所有照片识别结束。')


def updateFaceLib(group_id):
    '''
    更新人脸库
    '''
    createGroup(group_id)

    assess_token = baidu_ai_common.get_access_token()
    headers = {
        'Content-Type': 'application/json'
    }

    folder_path = "人脸库/"
    face_files = os.listdir(folder_path)

    for face_photo in face_files:
        face_photo_info = face_photo.split('.')
        face_photo_ext = face_photo_info[1]
        if face_photo_ext.lower() in ['jpg', 'jpeg', 'png', 'bmp']:
            user_info = face_photo_info[0].split('_')
            user_id = user_info[0]
            user_name = user_info[1]
            check_user_info = checkUserInfo(group_id, user_id)
            if check_user_info['error_code'] > 0:
                update_msg = '新增'
                url = "https://aip.baidubce.com/rest/2.0/face/v3/faceset/user/add?access_token=" + assess_token
                payload = json.dumps({
                    "group_id": group_id,
                    "image": baidu_ai_common.get_file_content_as_base64(folder_path+face_photo, False),
                    "image_type": "BASE64",
                    "user_id": user_id,
                    "user_info": user_name
                })
            else:
                update_msg = '更新'
                url = "https://aip.baidubce.com/rest/2.0/face/v3/faceset/user/update?access_token=" + assess_token
                payload = json.dumps({
                    "group_id": group_id,
                    "image": baidu_ai_common.get_file_content_as_base64(folder_path+face_photo, False),
                    "image_type": "BASE64",
                    "user_id": user_id,
                    "user_info": user_name
                })

            headers = {
                'Content-Type': 'application/json'
            }
            response = requests.request(
                "POST", url, headers=headers, data=payload)
            result = response.json()
            if result['error_code'] == 0:
                print('成功'+update_msg+'人脸信息：'+user_name)
            else:
                print('人脸信息更新失败：'+result['error_msg'])


def checkUserInfo(group_id, user_id):
    '''
    查询用户信息
    '''
    url = "https://aip.baidubce.com/rest/2.0/face/v3/faceset/user/get?access_token=" + \
        baidu_ai_common.get_access_token()

    payload = json.dumps({
        "group_id": group_id,
        "user_id": user_id
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)
    return response.json()


def createGroup(group_id):
    '''
    创建用户组，因如果用户组存在则只返回错误，并无其他问题，所以每次调用时都创建一下也无所谓
    '''
    url = "https://aip.baidubce.com/rest/2.0/face/v3/faceset/group/add?access_token=" + \
        baidu_ai_common.get_access_token()

    payload = json.dumps({
        "group_id": group_id  # 用户组id
    })
    headers = {
        'Content-Type': 'application/json'
    }
    requests.request("POST", url, headers=headers, data=payload)


if __name__ == '__main__':
    main()
